#include <opt_infer_engine.h>
#include <hilog/log.h>
#include <iostream>
#include <fstream>
#include <chrono>
namespace OPT {
    OptInferEngine::OptInferEngine(const std::string& model_path,
                                   const std::string &vocab_file,
                                   const std::string &merges_file)
    : processor_(vocab_file, merges_file)
    {
        loadModel(model_path);
    }

    void OptInferEngine::loadModel(const std::string& opt_model_path) {
        auto start = std::chrono::high_resolution_clock::now();
        try {
            opt_model_ = torch::jit::load(opt_model_path);
            opt_model_.eval();
        } catch (const c10::Error& e) {
            std::cerr << "error loading the model\n";
            std::cerr << e.what();
            OH_LOG_Print(LOG_APP, LOG_INFO, LOG_DOMAIN, "[OptNativeLoadModel]", "FAIL");
            exit(1);
        }
        auto end = std::chrono::high_resolution_clock::now();
        OH_LOG_Print(LOG_APP, LOG_INFO, LOG_DOMAIN, "[OptNativeLoadModel]", "DONE");
        std::cout << "OPT-125M Model is loaded in " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
    }

    std::string OptInferEngine::getReturnString(const std::string &text) {
        auto text_tensor = processor_.preprocessText(text);
        std::vector<torch::jit::IValue> inputs = {text_tensor};
        auto output = opt_model_.forward(inputs).toTensor(); // isTensorNow
        auto return_vec = processor_.getReturnString(output.reshape({-1}));
        std::string return_string;
        for (auto &str : return_vec) {
            return_string += str;
        }
        return return_string;
    }

    void OptInferEngine::test(const std::string &text){
        std::string checkString = "t t t";
        std::ofstream outfile(text + "/execution_times.txt");

        for(int i = 0; i < 1000; i++){
            checkString += "t ";
            auto start = std::chrono::high_resolution_clock::now();
            std::string return_string = getReturnString(checkString);
            auto end = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
            std::cout << return_string << std::endl;
            outfile << "Execution " << i + 1 << ": " << duration.count() << " microseconds" << std::endl;
        }
        outfile.close();
        return;
    }
}